# -*- coding: utf-8 *-*
import unittest
from graphs.listgraph import ListGraph
from graphs.matrixgraph import MatrixGraph


class PrimTest(unittest.TestCase):

    def test_prim_matrix(self):
        self.run_test1(MatrixGraph())
        self.run_test2(MatrixGraph())
        self.run_test3(MatrixGraph())

    def test_prim_list(self):
        self.run_test1(ListGraph())
        self.run_test2(ListGraph())
        self.run_test3(ListGraph())

    def run_test1(self, graph):
        graph.load('inputs/test1.txt')

        count = graph.vertex_count() - 1

        self.edges = graph.run_prim()

        self.assertEqual(count, self.edges.count())

        self.assertTrue(self.contains('s', 'u'))
        self.assertTrue(self.contains('u', 'x'))
        self.assertTrue(self.contains('x', 'y'))
        self.assertTrue(self.contains('x', 'z'))
        self.assertTrue(self.contains('s', 'v') or self.contains('v', 'x'))

    def run_test2(self, graph):
        graph.load('inputs/test2.txt')

        count = graph.vertex_count() - 1

        self.edges = graph.run_prim()

        self.assertEqual(count, self.edges.count())

        self.assertTrue(self.contains('s', 'u'))
        self.assertTrue(self.contains('u', 'v'))
        self.assertTrue(self.contains('v', 'x'))
        self.assertTrue(self.contains('x', 'y'))
        self.assertTrue(self.contains('y', 'z'))

    def run_test3(self, graph):
        graph.load('inputs/test3.txt')

        count = graph.vertex_count() - 1

        self.edges = graph.run_prim()

        self.assertEqual(count, self.edges.count())

        self.assertTrue(self.contains('a', 'b'))
        self.assertTrue(self.contains('a', 'c'))
        self.assertTrue(self.contains('c', 'd'))
        self.assertTrue(self.contains('d', 'e'))

    def contains(self, src, dest):
        for edge in self.edges:
            source = edge[0]
            destination = edge[1]
            if ((source == src and destination == dest) or
            (source == dest and destination == src)):
                return True
        return False
